package pipe

import "sync"

// 生产者，消费者模式消费
// 1. 生产者
// 2. 消费者
// 3. 两个chan，输入chan和输出chan
type InChan chan interface{}
type OutChan chan interface{}
type ProducerFunc func(in InChan)
type optionFunc func(*Pipe)

// 这个是消费函数
type CommonFunc func(item interface{}) interface{}

type Pipe struct {
	InChan   InChan
	Producer ProducerFunc
	Count    int //多少个消费者
}

func NewPipe() *Pipe {
	return &Pipe{
		Count: 1, //默认值
	}
}

// 设置数量
func WithPipeCount(count int) optionFunc {
	return func(p *Pipe) {
		p.Count = count
	}
}

func (this *Pipe) PipeDo(p ProducerFunc, c CommonFunc, options ...optionFunc) OutChan {
	//3. 设置数量
	if len(options) > 0 {
		for _, f := range options {
			f(this)
		}
	}
	//1. 执行生产者
	this.excProducer(p)
	//开始执行
	out := make(OutChan, 0)
	var wg = sync.WaitGroup{}
	for i := 1; i <= this.Count; i++ {
		wg.Add(1)
		go func(i InChan, f CommonFunc) {
			defer wg.Done()
			//开始消费，从管道里取数据出来进行消费
			for v := range i {
				out <- f(v)
			}
		}(this.InChan, c)
	}
	//返回数据
	go func() {
		defer close(out)
		wg.Wait()
	}()
	return out
}

// 执行生产者
func (this *Pipe) excProducer(p ProducerFunc) {
	this.InChan = make(InChan, 0)
	go func(inFunc ProducerFunc) {
		defer close(this.InChan)
		//需要在函数里面将数据塞进管道
		inFunc(this.InChan)
	}(p)
}
